#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import cv2
import torch
import numpy as np
import sys
import time
import os
from enum import Enum

# ROS2相关导入
import rclpy
from rclpy.node import Node
from geometry_msgs.msg import Twist
from sensor_msgs.msg import LaserScan
from std_msgs.msg import Bool

# 导入避障通用函数
import math
from yahboomcar_laser.common import *

os.environ['YOLO_VERBOSE'] = 'False'

RAD2DEG = 180 / math.pi

class FusionState(Enum):
    """融合控制状态枚举"""
    STOPPED = 0          # 停止
    FORWARD = 1          # 前进（缓速监控状态）
    TURN_LEFT = 5        # 左旋 (原地左转)
    TURN_RIGHT = 6       # 右旋 (原地右转)
    EMERGENCY_AVOIDANCE = 7  # 紧急避障

class BallTrackingLaserFusion(Node):
    def __init__(self, model_path, camera_device=0):
        """
        初始化融合控制器
        """
        super().__init__('ball_tracking_laser_fusion')
        
        self.model_path = model_path
        self.camera_device = camera_device
        
        # 初始化ROS2发布器和订阅器
        print("🔄 正在初始化ROS2控制...")
        try:
            self.pub_vel = self.create_publisher(Twist, '/cmd_vel', 1)
            self.pub_buzzer = self.create_publisher(Bool, '/Buzzer', 1)
            self.sub_laser = self.create_subscription(LaserScan, "/scan", self.laser_callback, 1)
            time.sleep(0.1)
            print("✅ ROS2控制初始化成功")
        except Exception as e:
            print(f"❌ ROS2控制初始化失败: {e}")
            raise
        
        # 加载YOLO模型
        print("🔄 正在加载绿色小球检测模型...")
        try:
            self.model = torch.hub.load('/root/yolov5-7.0', 'custom', 
                                      path=model_path, source='local', force_reload=True)
            self.model.conf = 0.4  # 置信度改为0.4
            self.model.iou = 0.45
            print(f"✅ 模型加载成功: {model_path}")
        except Exception as e:
            print(f"❌ 模型加载失败: {e}")
            raise
        
        # 初始化摄像头
        print(f"🔄 正在初始化摄像头...")
        self.cap = cv2.VideoCapture(camera_device)
        if not self.cap.isOpened():
            raise RuntimeError(f"❌ 无法打开摄像头: /dev/video{camera_device}")
        
        # 设置摄像头参数
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)
        self.cap.set(cv2.CAP_PROP_FPS, 30)
        
        self.frame_width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.frame_height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        print(f"✅ 摄像头初始化成功，分辨率: {self.frame_width}x{self.frame_height}")
        
        # ==================== 控制参数 ====================
        self.forward_speed = 0.2       # 前进速度 (m/s)
        self.turn_speed = 0.5          # 旋转速度 (rad/s)
        # ================================================
        
        # YOLO跟踪参数
        self.center_tolerance = 10.0  # 中心容忍度（百分比）
        self.min_ball_size = 5.0      # 最小球体尺寸（百分比）
        self.max_ball_size = 70.0     # 最大球体尺寸（百分比）
        
        # 激光避障参数
        self.LaserAngle = 40.0
        self.ResponseDist = 0.55  # 避障距离改为0.55米
        
        # YOLO吸引力因子
        self.yolo_attraction_factor = 6.0  # 基础吸引力因子
        
        # 障碍物检测
        self.Right_warning = 0
        self.Left_warning = 0
        self.front_warning = 0
        self.Right_threat_score = 0.0
        self.Left_threat_score = 0.0
        
        # 最新的激光数据
        self.last_laser_data = None
        
        # 状态机
        self.current_state = FusionState.STOPPED
        self.state_start_time = time.time()
        self.detection_count = 0
        
        # 紧急避障相关
        self.is_emergency_avoiding = False
        self.emergency_start_time = 0.0
        
        # ROS2控制器
        self.ros_ctrl = SinglePID()
        
    def laser_callback(self, scan_data):
        """激光雷达数据回调"""
        if not isinstance(scan_data, LaserScan): 
            return
        
        # 存储最新的激光数据
        self.last_laser_data = scan_data
        
    def process_laser_data_with_ball_mask(self, scan_data, ball_info):
        """
        处理激光雷达数据，根据小球位置屏蔽对应区域
        """
        # 重置计数器
        self.Right_warning = 0
        self.Left_warning = 0
        self.front_warning = 0
        self.Right_threat_score = 0.0
        self.Left_threat_score = 0.0
        
        ranges = np.array(scan_data.ranges)
        
        # 定义屏蔽区域：摄像头视野范围 -150° → -180° 和 +180° → +165°
        mask_ranges = []
        if ball_info is not None:
            # 当检测到小球时，屏蔽摄像头视野范围的激光数据
            mask_ranges = [(-180, -150), (165, 180)]
        
        processed_points = 0
        masked_points = 0
        
        for i in range(len(ranges)):
            angle = (scan_data.angle_min + scan_data.angle_increment * i) * RAD2DEG
            distance = ranges[i]
            
            # 过滤无效距离
            if distance <= 0.1 or distance > 10.0:
                continue
            
            processed_points += 1
            
            # 检查是否在屏蔽区域内
            should_mask = False
            for mask_start, mask_end in mask_ranges:
                if mask_start <= angle <= mask_end:
                    should_mask = True
                    masked_points += 1
                    break
            
            if should_mask:
                continue  # 跳过被屏蔽的数据
            
            # 右侧区域检测 (统一使用ResponseDist)
            if 160 > angle > 180 - self.LaserAngle:
                if distance < self.ResponseDist:
                    self.Right_warning += 1
                    self.Right_threat_score += 1.0 / distance
                    
            # 左侧区域检测
            elif -160 < angle < self.LaserAngle - 180:
                if distance < self.ResponseDist:
                    self.Left_warning += 1
                    self.Left_threat_score += 1.0 / distance
                    
            # 前方区域检测
            elif abs(angle) > 160:
                if distance <= self.ResponseDist:
                    self.front_warning += 1
        
        # 根据YOLO检测结果调整威胁度
        if ball_info is not None:
            center_x = ball_info['center_x']
            confidence = ball_info['confidence']
            ball_size = ball_info['ball_size']
            
            # 计算动态吸引力：基础因子 * (1+置信度) * (1+大小/100)
            dynamic_attraction = self.yolo_attraction_factor * (1 + confidence) * (1 + ball_size / 100.0)
            
            # 判断小球在左侧还是右侧
            if center_x < 50.0 - self.center_tolerance:  # 小球在左侧
                original_left = self.Left_threat_score
                self.Left_threat_score = max(0, self.Left_threat_score - dynamic_attraction)
            elif center_x > 50.0 + self.center_tolerance:  # 小球在右侧
                original_right = self.Right_threat_score
                self.Right_threat_score = max(0, self.Right_threat_score - dynamic_attraction)
        
    def detect_balls(self, frame):
        """
        检测当前帧中的绿色小球
        返回置信度最高的小球信息
        """
        results = self.model(frame)
        detections = results.pandas().xyxy[0]
        
        best_ball = None
        highest_confidence = 0.0
        
        for _, detection in detections.iterrows():
            confidence = detection['confidence']
            if confidence > 0.4 and confidence > highest_confidence:  # 置信度阈值改为0.4
                # 计算相对坐标
                xmin, ymin = detection['xmin'], detection['ymin']
                xmax, ymax = detection['xmax'], detection['ymax']
                
                center_x = ((xmin + xmax) / 2 / self.frame_width) * 100
                center_y = ((ymin + ymax) / 2 / self.frame_height) * 100
                width = ((xmax - xmin) / self.frame_width) * 100
                height = ((ymax - ymin) / self.frame_height) * 100
                
                # 检查球体大小是否在合理范围内
                ball_size = (width + height) / 2
                if self.min_ball_size <= ball_size <= self.max_ball_size:
                    best_ball = {
                        'confidence': confidence,
                        'center_x': center_x,
                        'center_y': center_y,
                        'width': width,
                        'height': height,
                        'ball_size': ball_size
                    }
                    highest_confidence = confidence
        
        return best_ball
    
    def decide_next_state(self, ball_info):
        """
        融合决策：根据小球位置和激光雷达数据决定下一个状态
        """
        if self.last_laser_data is None:
            return FusionState.STOPPED
        
        # 处理激光数据（考虑小球屏蔽）
        self.process_laser_data_with_ball_mask(self.last_laser_data, ball_info)
        
        # 检查是否需要紧急避障
        if self.front_warning > 10 and self.Left_warning > 10 and self.Right_warning > 10:
            return FusionState.EMERGENCY_AVOIDANCE
        
        # 统一使用ResponseDist阈值进行决策
        front_check = self.front_warning > 10
        left_check = self.Left_warning > 10
        right_check = self.Right_warning > 10
        
        # 获取威胁度（已经考虑了YOLO吸引力）
        left_threat = self.Left_threat_score
        right_threat = self.Right_threat_score
        
        # 新的决策逻辑：优先使用威胁度进行决策
        
        # 1. 检查紧急情况：前方+左右都有严重障碍
        if front_check and left_check and right_check:
            return FusionState.EMERGENCY_AVOIDANCE
        
        # 2. 前方有障碍：根据威胁度选择转向
        if front_check:
            if left_threat <= right_threat:
                return FusionState.TURN_LEFT
            else:
                return FusionState.TURN_RIGHT
        
        # 3. 前方无障碍：根据情况决策
        else:
            # 3.1 如果左右都有障碍，根据威胁度选择方向
            if left_check or right_check:
                if left_threat <= right_threat:
                    return FusionState.TURN_LEFT
                else:
                    return FusionState.TURN_RIGHT
            
            # 3.2 左右都无障碍：根据小球位置决策
            elif ball_info is not None:
                center_x = ball_info['center_x']
                ball_size = ball_info['ball_size']
                
                # 计算与画面中心的偏差
                center_offset = center_x - 50.0
                
                if abs(center_offset) <= self.center_tolerance:
                    # 小球在中心区域
                    if ball_size < 30.0:  # 球体较小，需要前进
                        return FusionState.FORWARD
                    else:  # 球体较大，停止
                        return FusionState.STOPPED
                elif center_offset < -self.center_tolerance:
                    # 小球在左侧，需要左旋
                    return FusionState.TURN_LEFT
                else:
                    # 小球在右侧，需要右旋
                    return FusionState.TURN_RIGHT
            
            # 3.3 无障碍且无小球：正常前进
            else:
                return FusionState.FORWARD
        
        return FusionState.FORWARD
    
    def publish_motion(self, linear_x=0.0, angular_z=0.0):
        """发布运动指令"""
        twist = Twist()
        twist.linear.x = linear_x
        twist.angular.z = angular_z
        self.pub_vel.publish(twist)
        
    def sound_buzzer(self, enable=True):
        """蜂鸣器控制"""
        buzzer_msg = Bool()
        buzzer_msg.data = enable
        self.pub_buzzer.publish(buzzer_msg)
    
    def execute_state(self, state):
        """
        执行状态对应的动作
        """
        if state == FusionState.STOPPED:
            self.publish_motion(0.0, 0.0)  # 停止
        elif state == FusionState.FORWARD:
            self.publish_motion(self.forward_speed, 0.0)  # 前进（缓速）
        elif state == FusionState.TURN_LEFT:
            self.publish_motion(0.0, self.turn_speed)  # 左旋（原地左转）
        elif state == FusionState.TURN_RIGHT:
            self.publish_motion(0.0, -self.turn_speed)  # 右旋（原地右转）
        elif state == FusionState.EMERGENCY_AVOIDANCE:
            self.publish_motion(0.0, self.turn_speed * 0.5)  # 紧急避障（慢速左转）
            self.sound_buzzer(True)
    
    def transition_to(self, new_state):
        """状态转换"""
        if new_state != self.current_state:
            self.current_state = new_state
            self.state_start_time = time.time()
            
            # 特殊状态处理
            if new_state == FusionState.EMERGENCY_AVOIDANCE:
                self.is_emergency_avoiding = True
                self.emergency_start_time = time.time()
            else:
                self.is_emergency_avoiding = False
                self.sound_buzzer(False)
    
    def get_state_duration(self):
        """获取当前状态持续时间"""
        return time.time() - self.state_start_time
    
    def state_machine_update(self, ball_info):
        """
        状态机更新
        """
        # 检查紧急避障超时
        if self.is_emergency_avoiding:
            if time.time() - self.emergency_start_time >= 1.0:
                self.is_emergency_avoiding = False
                self.sound_buzzer(False)
                # 紧急避障结束后转入正常状态
                next_state = self.decide_next_state(ball_info)
                self.transition_to(next_state)
                self.execute_state(next_state)
                return
            else:
                # 继续紧急避障
                self.execute_state(FusionState.EMERGENCY_AVOIDANCE)
                return
        
        # 检查转弯超时（3秒）
        if (self.current_state in [FusionState.TURN_LEFT, FusionState.TURN_RIGHT] 
            and self.get_state_duration() > 3.0):
            next_state = FusionState.FORWARD
        else:
            # 正常决策
            next_state = self.decide_next_state(ball_info)
        
        # 状态转换和执行
        self.transition_to(next_state)
        self.execute_state(next_state)
    
    def run(self):
        """
        主运行循环
        """
        print("\n🚀 开始融合控制...")
        print("📝 状态说明: STOPPED=停止, FORWARD=前进, TURN_LEFT=左旋, TURN_RIGHT=右旋, EMERGENCY_AVOIDANCE=紧急避障")
        print("🎯 融合策略: 视觉引导 + 激光避障")
        print("⌨️  按 Ctrl+C 退出\n")
        
        frame_count = 0
        last_fps_time = time.time()
        
        try:
            while rclpy.ok():
                # 读取摄像头帧
                ret, frame = self.cap.read()
                if not ret:
                    continue
                
                frame_count += 1
                
                # 检测绿色小球
                ball_info = self.detect_balls(frame)
                
                # 状态机更新（以YOLO频率为主）
                self.state_machine_update(ball_info)
                
                # 统计检测次数
                if ball_info:
                    self.detection_count += 1
                
                # 处理ROS2回调
                rclpy.spin_once(self, timeout_sec=0.001)
                
        except KeyboardInterrupt:
            print("\n🛑 用户停止程序")
        except Exception as e:
            print(f"\n💥 程序运行出错: {str(e)}")
        finally:
            self.cleanup()
    
    def cleanup(self):
        """
        清理资源
        """
        print("🧹 正在清理资源...")
        
        # 停止小车
        for _ in range(3):
            self.publish_motion(0.0, 0.0)  # 停止
            self.sound_buzzer(False)
            time.sleep(0.05)
        
        # 释放摄像头
        if hasattr(self, 'cap') and self.cap:
            self.cap.release()
        
        cv2.destroyAllWindows()
        print(f"✅ 资源清理完成，总共检测到 {self.detection_count} 次绿色小球")

def main():
    model_path = "/root/yolov5-7.0/model/best.pt"
    
    print("🤖 融合控制程序 - 视觉追踪 + 激光避障")
    print("=" * 60)
    print(f"📁 模型路径: {os.path.abspath(model_path)}")
    print(f"📁 工作目录: {os.getcwd()}")
    print("🎮 控制参数:")
    print("  - 前进速度: 0.2 m/s")
    print("  - 旋转速度: 0.5 rad/s")
    print("  - YOLO吸引力: 基础2.0 * (1+置信度) * (1+大小/100)")
    print("  - 避障距离阈值: 0.55m")
    print("  - YOLO置信度阈值: 0.4")
    print("📋 融合策略:")
    print("  - 摄像头视野屏蔽: -150°→-180°, +165°→+180°")
    print("  - 威胁度 = 激光威胁度 - YOLO吸引力")
    print("  - 频率同步: 以YOLO为主频率")
    
    # 初始化ROS2
    rclpy.init()
    
    try:
        fusion_controller = BallTrackingLaserFusion(model_path, camera_device=0)
        fusion_controller.run()
        
    except FileNotFoundError as e:
        print(f"📁 {str(e)}")
        print("💡 请确保模型文件 best.pt 在 model/ 目录中")
        return 1
    except Exception as e:
        print(f"💥 程序运行出错: {str(e)}")
        return 1
    finally:
        if rclpy.ok():
            rclpy.shutdown()
    
    return 0

if __name__ == "__main__":
    exit(main())